#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>

#include "common.h"
#include "cpu_info.h"

extern void start(CodeConfig *codeCfg);
extern void get_run_frequency(CodeConfig *codeCfg);

void print_usage()
{
    printf("\n");
    printf("nanobenchsuit usage:\n");
    printf("\n");
    printf(" -b, --binary <file>           Print the raw binary to the file.\n");
    printf(" -c, --core <n>                Pins the measurement thread to core n.\n");
    printf(" -e, --events <string>         Specify pmu grab events, use ',' to split.\n");
    printf(" -f, --freqency <n>            Test the running frequecy.\n");
    printf(" -i, --info <n>                Show the cpu info and pmu info. \n");
    printf(" -l, --loop <n>                Number of iterations of the inner loop.\n");
    printf(" -r, --run <n>                 Number of times the measurements are repeated.\n");
    printf(" -u, --unroll <n>              Number of copies of the benchmark code inside the inner loop.\n");
    return;
}

void print_code_cfg(CodeConfig *cfg)
{
    printf("loop        : %d\n"
           "unroll      : %d\n"
           "core        : %d\n"
           "run         : %d\n"
           "pmu_cnt     : %d\n"
           "event_cnt   : %d\n",
           cfg->loop, cfg->unroll, cfg->core, cfg->run, cfg->pmu_cnt, cfg->events_cnt);

    printf("binary_file : %s\n", cfg->binary_file != NULL ? cfg->binary_file : "null");
    printf("pmu_events  : [\n");
    for (int i = 0; i < cfg->events_cnt; i++) {
        printf("                 %s : %lu\n", cfg->events_list[i].name, cfg->events_list[i].value);
    }
    printf("              ]\n");
    return;
}

void free_code_cfg(CodeConfig *cfg)
{
    if (NULL != cfg->events_list) {
        free(cfg->events_list);
    }
    if (NULL != cfg->binary_file) {
        free(cfg->binary_file);
    }
    return;
}

PmuDataList *read_pmu_events(char *file_name, int *cnt)
{
    FILE *fd = fopen(file_name, "r");
    if (fd == NULL) {
        printf("open \"%s\" error!\n", file_name);
        printf("you need to run \"pmu_config.py\" with python3.\n");
        return NULL;
    }
    if (EOF == fscanf(fd, "%d", cnt)) {
        printf("read error\n");
        return NULL;
    }
    if (cnt <= 0) {
        return NULL;
    }

    PmuDataList *data = calloc(*cnt, sizeof(PmuDataList));
    if (NULL == data) {
        return NULL;
    }
    PmuDataList *tmp;
    for (int i = 0; i < *cnt; i++) {
        tmp = &data[i];
        if (EOF == fscanf(fd, "%s", tmp->name)) {
            printf("read name error\n");
            return NULL;
        }
        if (EOF == fscanf(fd, "%lx", &(tmp->data))) {
            printf("read data error\n");
            return NULL;
        }
    }
    fclose(fd);

    return data;
}

PmuDataList *get_capture_pmu_events(PmuDataList *pmuList, int pmuListCnt, char *arg, int *len)
{
    PmuDataList *capture_events = calloc(pmuListCnt, sizeof(PmuDataList));
    if (capture_events == NULL) {
        return NULL;
    }
    char *captured = calloc(pmuListCnt, sizeof(char));
    if (NULL == captured) {
        free(capture_events);
        return NULL;
    }

    int matched;
    const char *split = ",";
    PmuDataList *list;
    char *p;
    p = strtok(arg, split);
    while (p != NULL) {
        matched = 0;
        for (int i = 0; i < pmuListCnt; i++) {
            list = &(pmuList[i]);
            // check name
            if (0 != strcmp(p, list->name)) {
                continue;
            }
            matched = 1;
            // check duplication
            if (captured[i] == 1) {
                continue;
            }
            capture_events[*len].data = list->data;
            strcpy(capture_events[*len].name, list->name);
            *len += 1;
            captured[i] = 1;
            break;
        }
        // check invalid event name
        if (matched == 0) {
            printf("Not support event : %s\n", p);
            free(captured);
            free(capture_events);
            return NULL;
        }
        p = strtok(NULL, split);
    }
    free(captured);
    return capture_events;
}

int main(int argc, char **argv)
{
    struct option long_opts[] = { { "info", no_argument, 0, 'i' },
                                  { "binary", required_argument, 0, 'b' },
                                  { "run", required_argument, 0, 'r' },
                                  { "unroll", required_argument, 0, 'u' },
                                  { "loop", required_argument, 0, 'l' },
                                  { "core", required_argument, 0, 'c' },
                                  { "output", required_argument, 0, 'o' },
                                  { "events", required_argument, 0, 'e' },
                                  { "help", no_argument, 0, 'h' },
                                  { "frequency", no_argument, 0, 'f' },
                                  { 0, 0, 0, 0 } };

    CodeConfig codeCfg = {
        .loop = 1000,
        .unroll = 10,
        .core = -1,
        .run = 1,
        .pmu_cnt = 0,
        .events_cnt = 0,
        .events_list = NULL,
        .binary_file = NULL,
    };

    // default event :
    codeCfg.events_list = calloc(2, sizeof(PmuDataList));
    if (NULL == codeCfg.events_list) {
        return 1;
    }
    strcpy(codeCfg.events_list[0].name, "cycles");
    codeCfg.events_list[0].data = CYCLES;
    strcpy(codeCfg.events_list[1].name, "instructions");
    codeCfg.events_list[1].data = INSTRUCTIONS;
    codeCfg.events_cnt = 2;

    // read support pmu event list
    PmuDataList *pmuList = NULL;
    int pmuListCnt = 0;
    pmuList = read_pmu_events("pmu_events", &pmuListCnt);
    if (NULL == pmuList) {
        return 1;
    }

    int option = 0;
    while ((option = getopt_long_only(argc, argv, "", long_opts, NULL)) != -1) {
        switch (option) {
        case 0: break;
        case 'b':
            codeCfg.binary_file = calloc(strlen(optarg) + 1, sizeof(char));
            strcpy(codeCfg.binary_file, optarg);
            break;
        case 'c': codeCfg.core = atoi(optarg); break;
        case 'i':
            printf("\n");
            printf("+---------------------------------------------------------------+\n");
            printf("| cpu info :                                                    |\n");
            printf("+---------------------------------------------------------------+\n");
            get_cpu_info();
            printf("\n");
            printf("+---------------------------------------------------------------+\n");
            printf("| support pmu events :                                          |\n");
            printf("+---------------------------------------------------------------+\n");
            PmuDataList *p;
            for (int i = 0; i < pmuListCnt; i++) {
                p = &pmuList[i];
                printf("%s : %#lx\n", p->name, p->data);
            }
            return 0;
        case 'r': codeCfg.run = atoi(optarg); break;
        case 'f': get_run_frequency(&codeCfg); return 0;
        case 'u':
            codeCfg.unroll = atoi(optarg);
            if (codeCfg.unroll <= 0) {
                printf("Error: unroll_count must be > 0\n");
                return 1;
            }
            break;
        case 'l': codeCfg.loop = atoi(optarg); break;
        case 'e':
            codeCfg.events_cnt = 0;
            free(codeCfg.events_list);
            codeCfg.events_list = NULL;
            codeCfg.events_list = get_capture_pmu_events(pmuList, pmuListCnt, optarg, &codeCfg.events_cnt);
            if (NULL == codeCfg.events_list) {
                return 1;
            }
            break;
        default: print_usage(); return 1;
        }
    }

    // check pmu version
    if (2 > get_pmu_verison()) {
        printf("pmu version too low.\n");
        return 1;
    }

    // check pmu counters
    codeCfg.pmu_cnt = get_pmu_counters();
    if (codeCfg.pmu_cnt <= 0) {
        printf("pmu event counters not support.\n");
        return 1;
    }

    start(&codeCfg);
    print_code_cfg(&codeCfg);

    free(pmuList);
    free_code_cfg(&codeCfg);

    return 0;
}
